{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 库的导入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision.transforms.functional import to_tensor, to_pil_image\n",
    "\n",
    "from captcha.image import ImageCaptcha\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "import string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0123456789 80 40 10\n"
     ]
    }
   ],
   "source": [
    "characters =  string.digits\n",
    "width, height, n_len, n_classes = 80, 40, 4, len(characters)\n",
    "label_length = 2\n",
    "# n_input_length = 12\n",
    "print(characters, width, height, n_classes)\n",
    "\n",
    "max_epoch = 15\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "init net\n",
      "epoch:0  train_loss:0.1502\n",
      "test_acc: 0.4347\n",
      "epoch:1  train_loss:0.0868\n",
      "test_acc: 0.6752\n",
      "epoch:2  train_loss:0.0616\n",
      "test_acc: 0.7367\n",
      "epoch:3  train_loss:0.0591\n",
      "test_acc: 0.701\n",
      "epoch:4  train_loss:0.0431\n",
      "test_acc: 0.8632\n",
      "epoch:5  train_loss:0.0518\n",
      "test_acc: 0.8946\n",
      "epoch:6  train_loss:0.0549\n",
      "test_acc: 0.3348\n",
      "epoch:7  train_loss:0.0429\n",
      "test_acc: 0.8665\n",
      "epoch:8  train_loss:0.0272\n",
      "test_acc: 0.8717\n",
      "epoch:9  train_loss:0.0373\n",
      "test_acc: 0.8864\n",
      "epoch:10  train_loss:0.0242\n",
      "test_acc: 0.884\n",
      "epoch:11  train_loss:0.0287\n",
      "test_acc: 0.9302\n",
      "epoch:12  train_loss:0.0366\n",
      "test_acc: 0.9051\n",
      "epoch:13  train_loss:0.0245\n",
      "test_acc: 0.9254\n",
      "epoch:14  train_loss:0.0276\n",
      "test_acc: 0.8764\n"
     ]
    }
   ],
   "source": [
    "class TwoNumCaptchaDataset(Dataset):\n",
    "    \"\"\" 两位数字验证码的图片 \"\"\"\n",
    "    def __init__(self, characters, length, width, height, label_length):\n",
    "        super(TwoNumCaptchaDataset, self).__init__()\n",
    "        self.characters = characters\n",
    "        self.length = length\n",
    "        self.width = width\n",
    "        self.height = height\n",
    "        # self.input_length = input_length\n",
    "        self.label_length = label_length\n",
    "        self.n_class = len(characters)\n",
    "        self.generator = ImageCaptcha(width=width, height=height)\n",
    "\n",
    "    def encode_label(self,target_str):\n",
    "        target = []\n",
    "        for char in target_str:\n",
    "            vec = [0.0] * len(self.characters)\n",
    "            vec[self.characters.find(char)] = 1.0\n",
    "            target += vec\n",
    "        return target\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.length\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        random_str = ''.join([random.choice(self.characters[1:]) for j in range(self.label_length)])\n",
    "        image = to_tensor(self.generator.generate_image(random_str))\n",
    "        # 进行one_hot编码\n",
    "        target = torch.tensor(self.encode_label(random_str)) \n",
    "        # target_length = torch.full(size=(1, ), fill_value=self.label_length, dtype=torch.long)\n",
    "        return image, target\n",
    "\n",
    "\n",
    "class CNN(nn.Module):\n",
    "    def __init__(self,num_class=10, num_char=2):\n",
    "        super(CNN,self).__init__()\n",
    "        self.num_class = num_class\n",
    "        self.num_char = num_char\n",
    "        self.conv = nn.Sequential(\n",
    "                # b*3*80*40\n",
    "                nn.Conv2d(3, 16, 3, padding=(1, 1)),\n",
    "                nn.MaxPool2d(2, 2),\n",
    "                nn.BatchNorm2d(16),\n",
    "                nn.ReLU(),\n",
    "                \n",
    "                nn.Conv2d(16, 64, 3, padding=(1, 1)),\n",
    "                nn.MaxPool2d(2, 2),\n",
    "                nn.BatchNorm2d(64),\n",
    "                nn.ReLU(),\n",
    "                )\n",
    "        self.fc = nn.Linear(64*10*20, self.num_class*self.num_char)\n",
    "        # temp = torch.randn(2,3,40,80)\n",
    "        # out = self.conv(temp)\n",
    "        # print(out.shape)\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv(x)\n",
    "        # -1 适应后面的size\n",
    "        x = x.view(-1,64*10*20)\n",
    "        x = self.fc(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "def calculat_acc(output, target):\n",
    "    \"\"\" 计算准确率 \"\"\"\n",
    "    output, target = output.view(-1, 10), target.view(-1, 10)\n",
    "    output = nn.functional.softmax(output, dim=1)\n",
    "    output = torch.argmax(output, dim=1)\n",
    "    target = torch.argmax(target, dim=1)\n",
    "    output, target = output.view(-1, 2), target.view(-1, 2)\n",
    "    correct_list = []\n",
    "    for i, j in zip(target, output):\n",
    "        if torch.equal(i, j):\n",
    "            correct_list.append(1)\n",
    "        else:\n",
    "            correct_list.append(0)\n",
    "    acc = sum(correct_list) / len(correct_list)\n",
    "    return acc\n",
    "\n",
    "\n",
    "batch_size = 128\n",
    "train_set = TwoNumCaptchaDataset(\n",
    "    characters=characters, length=batch_size*100, width=width, height=height, label_length=2)\n",
    "train_loader = DataLoader(train_set, batch_size=batch_size,num_workers=12)\n",
    "test_set = TwoNumCaptchaDataset(\n",
    "    characters=characters, length=300, width=width, height=height, label_length=2)\n",
    "test_loader = DataLoader(test_set, batch_size=batch_size)\n",
    "\n",
    "cnn = CNN()\n",
    "print('init net')\n",
    "if torch.cuda.is_available():\n",
    "    cnn.cuda()\n",
    "optimizer = torch.optim.Adam(cnn.parameters(), lr=0.001)\n",
    "criterion = nn.MultiLabelSoftMarginLoss()\n",
    "# criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "for epoch in range(max_epoch):\n",
    "    cnn.train()\n",
    "    for img, target in train_loader:\n",
    "        if torch.cuda.is_available():\n",
    "            img = img.cuda()\n",
    "            target = target.cuda()\n",
    "        # print('target_shape',target.shape)\n",
    "        output = cnn(img)\n",
    "        # print('out_shape',output.shape)\n",
    "        loss = criterion(output, target)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    print(f'epoch:{epoch}  train_loss:{loss:.4f}')\n",
    "    \n",
    "    \n",
    "    # test\n",
    "    with torch.no_grad():\n",
    "        # loss_history = []\n",
    "        acc_history = []\n",
    "        cnn.eval()\n",
    "        for img, target in test_loader:\n",
    "            if torch.cuda.is_available():\n",
    "                img = img.cuda()\n",
    "                target = target.cuda()\n",
    "            test_out = cnn(img)\n",
    "            acc = calculat_acc(test_out, target)\n",
    "            acc_history.append(float(acc))\n",
    "            # loss_history.append(float(loss))\n",
    "        print('test_acc: {:.4}'.format(\n",
    "            # torch.mean(torch.Tensor(loss_history)),\n",
    "            torch.mean(torch.Tensor(acc_history))\n",
    "            ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch:0  train_loss:0.0259\n",
      "test_acc: 0.9688\n",
      "epoch:1  train_loss:0.0169\n",
      "test_acc: 0.946\n",
      "epoch:2  train_loss:0.0187\n",
      "test_acc: 0.9635\n",
      "epoch:3  train_loss:0.0115\n",
      "test_acc: 0.9612\n",
      "epoch:4  train_loss:0.0302\n",
      "test_acc: 0.9564\n",
      "epoch:5  train_loss:0.0200\n",
      "test_acc: 0.9512\n",
      "epoch:6  train_loss:0.0290\n",
      "test_acc: 0.9616\n",
      "epoch:7  train_loss:0.0166\n",
      "test_acc: 0.9583\n",
      "epoch:8  train_loss:0.0193\n",
      "test_acc: 0.9844\n",
      "epoch:9  train_loss:0.0152\n",
      "test_acc: 0.9484\n",
      "done\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# 降低学习效率继续学\n",
    "optimizer = torch.optim.Adam(cnn.parameters(), lr=0.0001)\n",
    "# criterion = nn.MultiLabelSoftMarginLoss()\n",
    "# criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "for epoch in range(10):\n",
    "    cnn.train()\n",
    "    for img, target in train_loader:\n",
    "        if torch.cuda.is_available():\n",
    "            img = img.cuda()\n",
    "            target = target.cuda()\n",
    "        # print('target_shape',target.shape)\n",
    "        output = cnn(img)\n",
    "        # print('out_shape',output.shape)\n",
    "        loss = criterion(output, target)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    print(f'epoch:{epoch}  train_loss:{loss:.4f}')\n",
    "    \n",
    "    \n",
    "    # test\n",
    "    with torch.no_grad():\n",
    "        # loss_history = []\n",
    "        acc_history = []\n",
    "        cnn.eval()\n",
    "        for img, target in test_loader:\n",
    "            if torch.cuda.is_available():\n",
    "                img = img.cuda()\n",
    "                target = target.cuda()\n",
    "            test_out = cnn(img)\n",
    "            acc = calculat_acc(test_out, target)\n",
    "            acc_history.append(float(acc))\n",
    "            # loss_history.append(float(loss))\n",
    "        print('test_acc: {:.4}'.format(\n",
    "            # torch.mean(torch.Tensor(loss_history)),\n",
    "            torch.mean(torch.Tensor(acc_history))\n",
    "            ))\n",
    "print('done')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "saved\n"
     ]
    }
   ],
   "source": [
    "torch.save(cnn.state_dict(), 'cnn_6_27.pth')\n",
    "print('saved')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
